use crate::controller::ControllerSet;
use crate::http::server::error::{Error, HttpResponseError};
use axum::Json;
use axum::body::Body;
use axum::extract::rejection::JsonRejection;
use axum::extract::{FromRequest, Request, State};
use axum::http::{Response, header};
use axum::middleware::Next;
use serde::de::DeserializeOwned;
use std::collections::HashMap;
use std::sync::Arc;
use validator::Validate;
use wanmesh_common::auth;

pub struct ValidatedJson<J>(pub J);

impl<T, S> FromRequest<S> for ValidatedJson<T>
where
    T: DeserializeOwned + Validate,
    S: Send + Sync,
    Json<T>: FromRequest<S, Rejection = JsonRejection>,
{
    type Rejection = Error;

    async fn from_request(req: Request, state: &S) -> Result<Self, Self::Rejection> {
        let Json(value) = Json::<T>::from_request(req, state).await?;
        value
            .validate()
            .map_err(|e| Error::ValidationErrors(e.into()))?;
        Ok(Self(value))
    }
}

pub struct AuthMiddleware {
    controller_set: Arc<ControllerSet>,
}

impl AuthMiddleware {
    pub fn new(controller_set: Arc<ControllerSet>) -> Self {
        Self { controller_set }
    }

    pub async fn auth(
        State(state): State<Arc<Self>>,
        request: Request,
        next: Next,
    ) -> Result<Response<Body>, Error> {
        let authorization = get_authorization(&request)?;
        if authorization.is_empty() {
            return Err(Error::HttpResponseError(HttpResponseError::unauthorized(
                "缺少Authorization".to_string(),
            )));
        }

        if !authorization.starts_with(auth::BEARER_PREFIX) {
            return Err(Error::HttpResponseError(HttpResponseError::unauthorized(
                "Authorization格式错误".to_string(),
            )));
        }

        let access_token = match authorization.strip_prefix(auth::BEARER_PREFIX) {
            None => {
                return Err(Error::HttpResponseError(HttpResponseError::unauthorized(
                    "缺少access_token".to_string(),
                )));
            }
            Some(v) => v.to_string(),
        };

        state
            .controller_set
            .auth
            .get_auth_info(access_token.as_str())
            .await?;

        Ok(next.run(request).await)
    }
}

fn get_authorization(request: &Request) -> Result<String, Error> {
    let connection = get_request_header(request, header::CONNECTION);
    let upgrade = get_request_header(request, header::UPGRADE);

    if connection.to_lowercase() == "upgrade" && upgrade == "websocket" {
        return get_authorization_ws(request);
    }

    Ok(get_request_header(request, header::AUTHORIZATION))
}

fn get_authorization_ws(request: &Request) -> Result<String, Error> {
    let query_params: axum::extract::Query<HashMap<String, String>> =
        axum::extract::Query::try_from_uri(request.uri()).map_err(anyhow::Error::new)?;

    let value = query_params
        .get(auth::AUTHORIZATION)
        .map_or("".to_string(), |v| v.clone());
    Ok(value)
}

fn get_request_header<K>(request: &Request, key: K) -> String
where
    K: header::AsHeaderName,
{
    request.headers().get(key).map_or("".to_string(), |v| {
        String::from_utf8_lossy(v.as_bytes()).to_string()
    })
}
